/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.is.cef.command.listener.exe;

import bd.org.apache.commons.math.stat.descriptive.SummaryStatistics;
import cef.command.Command;
import cef.command.listener.exe.CommandExeEvent;
import cef.command.listener.exe.FigureEditorExeComponent;
import coalescent.is.Samplers_K69;
import coalescent.is.cef.command.RunSamplers_Cmd;
import coalescent.is.cef.event.RunSamplers_Evt;
import org.apfloat.Apcomplex;
import org.apfloat.Apfloat;
import org.apfloat.ApfloatMath;
import org.jfree.chart.ChartPanel;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.axis.NumberAxis;
import org.jfree.chart.plot.ValueMarker;
import org.jfree.chart.plot.XYPlot;
import org.jfree.chart.renderer.xy.XYLineAndShapeRenderer;
import org.jfree.data.general.DatasetChangeEvent;
import org.jfree.data.xy.DefaultXYDataset;
import org.jfree.ui.LengthAdjustmentType;
import org.jfree.ui.RectangleAnchor;
import org.jfree.ui.TextAnchor;

import javax.swing.*;
import java.awt.*;
import java.math.BigDecimal;
import java.text.DecimalFormat;
import java.util.*;
import java.util.List;

import static commons.util.ApfloatUtil.Max;
import static commons.util.ApfloatUtil.Min;
import static commons.util.FormatUtil.format;

/**
 * @author Susanta Tewari
 * @version version
 * @history Created on Aug 3, 2012.
 * @since since
 */
public abstract class VizSamplers_EC extends FigureEditorExeComponent {

    /** formatting */
    private static final Apfloat LOG_BASE  = new Apfloat(10);
    private static final DecimalFormat fmt = new DecimalFormat("0.####");

    /** event cache */
    protected Collection<String> samplers;

    /** sampler maps & data set */
    private Map<String, List<Apfloat>> estimate_map;
    private Map<String, JLabel> itr_map;
    private Map<String, JLabel> measure_map;

    /** GUI settings */
    protected RunSamplers_Evt.GUI guiSettings;
    protected Map<String, Color> colorMap;
    protected int WINDOW_SIZE          = 15;    // sliding window
    private static final Stroke stroke = new BasicStroke(4.0f);

    /** chart */
    private Map<String, double[][]> data_map;
    protected DefaultXYDataset data_set;
    protected Apfloat rangeMin = Apcomplex.ZERO,
                      rangeMax = Apcomplex.ZERO;    // dynamically adjusted
    protected JFreeChart chart;

    /** update-1 controls */
    protected final JPanel update1_panel      = new JPanel();
    private static final GridBagConstraints c = new GridBagConstraints();

    /** update-2 controls */
    private int timer2Counter;

    @Override
    public void initComponents(final Command cmd) {

        final RunSamplers_Cmd command = (RunSamplers_Cmd) cmd;

        guiSettings = command.getEvent().getGUISettings();

        if (guiSettings.isISWindowSizeKnown()) WINDOW_SIZE = guiSettings.getISWindowSize();


        // sampler maps & data set
        colorMap     = guiSettings.getColorMap();
        samplers     = command.getSamplerNames();
        estimate_map = new HashMap<>(samplers.size());
        itr_map      = new HashMap<>(samplers.size());
        measure_map  = new HashMap<>(samplers.size());
        data_map     = new HashMap<>(samplers.size());
        data_set     = new DefaultXYDataset();

        for (final String samplerName : samplers) {

            estimate_map.put(samplerName, new LinkedList<Apfloat>());

            final JLabel jLabel = of_label();

            jLabel.setForeground(colorMap.get(samplerName));
            itr_map.put(samplerName, jLabel);

            final double[][] data = create_data();

            data_map.put(samplerName, data);
            data_set.addSeries(samplerName, data);
        }


        // init chart
        title = getTitle();

        final NumberAxis x_axis = new NumberAxis("Sampler observations (interpolated) over time");
        final NumberAxis y_axis = new NumberAxis(getYAxisTitle());

        x_axis.setAutoRangeIncludesZero(false);
        y_axis.setAutoRangeIncludesZero(false);

        final XYLineAndShapeRenderer renderer = new XYLineAndShapeRenderer(true, false);
        final XYPlot xyplot                   = new XYPlot(data_set, x_axis, y_axis, renderer);

        for (final String s : samplers) {

            final int series = data_set.indexOf(s);

            renderer.setSeriesPaint(series, colorMap.get(s));
            renderer.setSeriesStroke(series, stroke);
        }

        chart  = new JFreeChart(title, JFreeChart.DEFAULT_TITLE_FONT, xyplot, true);
        figure = new JPanel(new BorderLayout());

        figure.add(new ChartPanel(chart), BorderLayout.CENTER);


        // timer1 panel
        update1_panel.setLayout(new GridBagLayout());

        c.insets = new Insets(5, 5, 5, 5);
        c.gridx  = 0;
        c.gridy  = 0;


        // iterations
        update1_panel.add(new JLabel("Iterations: "), c);

        for (final String s : samplers) {

            c.gridy++;

            update1_panel.add(itr_map.get(s), c);
        }


        // measure
        c.gridx = 1;
        c.gridy = 0;

        update1_panel.add(new JLabel(getMeasureTitle() + ": "), c);

        for (final String samplerName : samplers) {

            c.gridy++;

            final JLabel jLabel = of_label();

            jLabel.setForeground(colorMap.get(samplerName));
            measure_map.put(samplerName, jLabel);
            update1_panel.add(jLabel, c);
        }

        figure.add(update1_panel, BorderLayout.EAST);
    }

    @Override
    public void updateComponentData(final CommandExeEvent event) {

        final RunSamplers_Cmd command = (RunSamplers_Cmd) event.getSource();

        update_1(command);

        if (timer2Counter % 3 == 0) update_2(command);

        timer2Counter = (timer2Counter >= 100) ? 0 : timer2Counter + 1;
    }

    @Override
    public void lastUpdate(final CommandExeEvent event) {

        final RunSamplers_Cmd command = (RunSamplers_Cmd) event.getSource();

        update_1(command);

        for (int i = 0; i < 3; i++) {
            update_2(command);    // plot requires min 3 points to display
        }
    }

    /**
     * Updates iteration count and IS measure.
     *
     * @param cmd command
     */
    protected void update_1(final RunSamplers_Cmd cmd) {

        for (final String sampler : samplers) {

            final long itrCount = cmd.getSampleSize(sampler);

            if (itrCount < 3) return;

            final JLabel jLabel = itr_map.get(sampler);

            jLabel.setText("<html>" + "<big>[" + format(itrCount) + "]</big></html>");
        }

        for (final String sampler : samplers) {

            final Apfloat estimate = getISMeasure(cmd, sampler);
            final JLabel jLabel    = measure_map.get(sampler);

            jLabel.setText("<html>" + "<big>[" + fmt.format(estimate.doubleValue())
                           + "]</big></html>");
        }
    }

    /**
     * updates graph.
     *
     * @param cmd command
     */
    protected void update_2(final RunSamplers_Cmd cmd) {

        final SummaryStatistics stat = new SummaryStatistics();

        rangeMin = Apcomplex.ZERO;
        rangeMax = Apcomplex.ZERO;

        for (final String sampler : samplers) {

            final List<Apfloat> buffer = estimate_map.get(sampler);
            final Apfloat estimate     = getISMeasure(cmd, sampler);

            if (estimate == null) break;    // cases for log zero etc..

            buffer.add(estimate);

            if (buffer.size() > WINDOW_SIZE) buffer.remove(0);


            // add data to data set
            double[][] data = data_map.get(sampler);

            for (int i = 0; i < buffer.size(); i++) {

                final Apfloat value = buffer.get(i);

                data[0][i] = i;
                data[1][i] = value.doubleValue();

                stat.addValue(data[1][i]);
            }


            // range min,max
            final Apfloat min = new Apfloat(stat.getMin());
            final Apfloat max = new Apfloat(stat.getMax());

            rangeMin = (rangeMin == Apcomplex.ZERO) ? min : Min(rangeMin, min);
            rangeMax = (rangeMax == Apcomplex.ZERO) ? max : Max(rangeMax, max);

            data_set.addSeries(sampler, data);
        }

        process_measure_specific(cmd, data_set);


        // adjust range axis
        final double se = stat.getStandardDeviation().doubleValue();

        chart.getXYPlot().getRangeAxis().setRange(rangeMin.doubleValue() - se,
                rangeMax.doubleValue() + se);


        // set the chart data set
        chart.getXYPlot().datasetChanged(new DatasetChangeEvent(this, data_set));
    }

    private static JLabel of_label() {

        final JLabel jLabel = new JLabel();

        jLabel.setOpaque(true);
        jLabel.setBackground(Color.white);

        return jLabel;
    }

    private static JLabel of_label(String text, Color foreColor) {

        final JLabel jLabel = new JLabel(text);

        jLabel.setOpaque(true);
        jLabel.setBackground(Color.white);
        jLabel.setForeground(foreColor);

        return jLabel;
    }

    private double[][] create_data() {

        final double[][] data = new double[2][WINDOW_SIZE];

        Arrays.fill(data[0], Double.NaN);
        Arrays.fill(data[1], Double.NaN);

        return data;
    }

    protected void process_measure_specific(final RunSamplers_Cmd command,
            final DefaultXYDataset dataset) {}

    protected abstract String getTitle();

    protected abstract String getMeasureTitle();

    protected abstract String getYAxisTitle();

    protected abstract Apfloat getISMeasure(RunSamplers_Cmd command, String samplerName);

    @Override
    public Class<? extends Command> getTargetCommandClass() {
        return RunSamplers_Cmd.class;
    }

    /**
     * IS ESS tracker
     */
    public static final class IS_TRACKER_ESS extends VizSamplers_EC {

        @Override
        protected Apfloat getISMeasure(final RunSamplers_Cmd command, final String samplerName) {

            final BigDecimal eSS = command.getIS_ESS(samplerName);

            return new Apfloat(eSS);
        }

        @Override
        protected String getMeasureTitle() {
            return "ESS";
        }

        @Override
        protected String getTitle() {
            return "Samplers on measure: Effective Sample Size (ESS)";
        }

        @Override
        protected String getYAxisTitle() {
            return "Effective Sample Size";
        }

        @Override
        public IS_TRACKER_ESS newInstance() {
            return new IS_TRACKER_ESS();
        }
    }

    /**
     * IS Mean tracker
     */
    public static final class IS_TRACKER_MEAN extends VizSamplers_EC {

        private Apfloat logExactProb;
        private boolean exactValueKnown;
        private Map<String, JLabel> mean_labels_map = new HashMap<>();

        @Override
        public void initComponents(Command cmd) {

            super.initComponents(cmd);

            exactValueKnown = guiSettings.isExactValueKnown();

            if (exactValueKnown) {

                final BigDecimal prob   = guiSettings.getExactValue();
                final Color exact_color = colorMap.get(Samplers_K69.EXACT);

                logExactProb = ApfloatMath.log(new Apfloat(prob), LOG_BASE);


                // add marker
                ValueMarker valuemarker = new ValueMarker(logExactProb.doubleValue());

                valuemarker.setLabelOffsetType(LengthAdjustmentType.EXPAND);
                valuemarker.setPaint(exact_color);
                valuemarker.setStroke(new BasicStroke(2.0f));
                valuemarker.setLabel("Exact Probability");
                valuemarker.setLabelAnchor(RectangleAnchor.BOTTOM_RIGHT);
                valuemarker.setLabelTextAnchor(TextAnchor.TOP_RIGHT);
                chart.getXYPlot().addRangeMarker(valuemarker);


                // add label to means panel
                final String prob_s       = fmt.format(logExactProb.doubleValue());
                final JLabel exact_jLabel = of_label(format(prob_s), exact_color);

                c.gridy++;

                update1_panel.add(exact_jLabel, c);

                c.gridx = 0;

                update1_panel.add(of_label(format("Exact"), exact_color), c);
            }
        }

        private static String format(String val) {

            return //J-
                    "<html>" +
                       "<big>" +
                        "[" + val + "]" +
                       "</big>" +
                    "</html>"; //J+

        }

        @Override
        protected String getMeasureTitle() {
            return "Means";
        }

        @Override
        protected Apfloat getISMeasure(final RunSamplers_Cmd command, final String samplerName) {

            final BigDecimal mean = command.getISMean(samplerName);

            if (mean.equals(BigDecimal.ZERO)) {
                return null;
            }

            return ApfloatMath.log(new Apfloat(mean), LOG_BASE);
        }

        @Override
        protected void process_measure_specific(final RunSamplers_Cmd command,
                final DefaultXYDataset dataset) {

            if (exactValueKnown) {

                rangeMin = Min(rangeMin, logExactProb);
                rangeMax = Max(rangeMax, logExactProb);
            }
        }

        @Override
        protected String getTitle() {
            return "Samplers on measure: Mean";
        }

        @Override
        protected String getYAxisTitle() {
            return "Mean (Lg10)";
        }

        @Override
        public IS_TRACKER_MEAN newInstance() {
            return new IS_TRACKER_MEAN();
        }
    }

    /**
     * IS SE tracker
     */
    public static final class IS_TRACKER_SE extends VizSamplers_EC {

        @Override
        protected String getMeasureTitle() {
            return "Std. Error";
        }

        @Override
        protected Apfloat getISMeasure(final RunSamplers_Cmd command, final String samplerName) {

            final BigDecimal se = command.getIS_SE(samplerName);

            if (se.equals(BigDecimal.ZERO)) {
                return null;
            }

            return ApfloatMath.log(new Apfloat(se), LOG_BASE);
        }

        @Override
        protected String getTitle() {
            return "Samplers on measure: Std. Error";
        }

        @Override
        protected String getYAxisTitle() {
            return "Standard Error (Lg10)";
        }

        @Override
        public IS_TRACKER_SE newInstance() {
            return new IS_TRACKER_SE();
        }
    }
}
